import math
import os
import json
from PIL import Image
import shutil


# 定义类别和对应的group_id
category_to_group_id = {
    "box": 0,
    "arrows_right": 1,
    "arrows_left": 2
}


def convert_to_yolov8_pose(json_file_path, output_file_path):
    with open(json_file_path, 'r') as file:
        data = json.load(file)

    image_path = json_file_path.replace(".json", ".jpg")

    # 尝试打开图片文件，若文件不存在则跳过
    try:
        image = Image.open(image_path)
        image_width, image_height = image.size
    except FileNotFoundError:
        print(f"File not found: {image_path}. Skipping this file.")
        return

    is_wrong = False
    for shape in data["shapes"]:
        if (
                # shape["label"] == "box"
                # or
                shape["label"] == "arrows_right"
                # or
                # shape["label"] == "arrows_left"
        ):
           is_wrong = True
        if shape["shape_type"] == "rectangle":
            for keypoint_label in [f"{shape['label']}_{i}" for i in range(1, 7)]:
                keypoint = next((keypoint for keypoint in data["shapes"] if keypoint["label"] == keypoint_label), None)
                if keypoint:
                    keypoint_x = keypoint["points"][0][0]
                    keypoint_y = keypoint["points"][0][1]
                    if math.sqrt(keypoint_x**2 + keypoint_y**2) < 0.5:
                        is_wrong = True
                else:
                    is_wrong = True

    if is_wrong:
        # 获取文件名
        file_name = os.path.basename(json_file_path)
        # 构造目标路径
        destination_file = os.path.join(output_folder, file_name)
        # 移动文件
        shutil.move(json_file_path, destination_file)
        # 获取文件名
        file_name = os.path.basename(image_path)
        # 构造目标路径
        destination_file = os.path.join(output_folder, file_name)
        # 移动文件
        shutil.move(image_path, destination_file)

# 指定存储标注文件和图像文件的文件夹路径
folder_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/good restore (copy)/frame_and_side/left"
# output_folder = os.path.join(os.path.dirname(folder_path), "wrong")
output_folder = os.path.join(folder_path, "wrong")

os.makedirs(output_folder, exist_ok=True)

for filename in os.listdir(folder_path):
    if filename.endswith(".json"):
        json_file_path = os.path.join(folder_path, filename)
        # output_file_path = os.path.join(output_folder, filename.replace(".json", ".txt"))

        convert_to_yolov8_pose(json_file_path, output_folder)

print("Conversion completed.")

